// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <melon/container/f14_map.h>

#include <pollux/expression/eval_ctx.h>
#include <pollux/expression/expr.h>
#include <pollux/expression/vector_function.h>
#include <pollux/functions/lib/comparator_util.h>
#include <pollux/functions/lib/rows_translation_util.h>

namespace kumo::pollux::functions {
    namespace {
        // See documentation at https://prestodb.io/docs/current/functions/array.html
        ///
        /// Implements the array_duplicates function.
        ///
        /// Along with the hash map, we maintain a `hasNull` flag that indicates
        /// whether null is present in the array.
        ///
        /// Zero element copy:
        ///
        /// In order to prevent copies of array elements, the function reuses the
        /// internal elements() vector from the original ArrayVector.
        ///
        /// First a new vector is created containing the indices of the elements
        /// which will be present in the output, and wrapped into a DictionaryVector.
        /// Next the `sizes` and `offsets` vectors that control where output arrays
        /// start and end are wrapped into the output ArrayVector.
        template<typename T>
        class ArrayDuplicatesFunction : public exec::VectorFunction {
        public:
            void apply(
                const SelectivityVector &rows,
                std::vector<VectorPtr> &args, // Not using const ref so we can reuse args
                const TypePtr &outputType,
                exec::EvalCtx &context,
                VectorPtr &result) const override {
                auto &arg = args[0];

                VectorPtr localResult;

                // Input can be constant or flat.
                if (arg->is_constant_encoding()) {
                    auto *constantArray = arg->as<ConstantVector<ComplexType> >();
                    const auto &flatArray = constantArray->value_vector();
                    const auto flatIndex = constantArray->index();

                    exec::LocalSingleRow singleRow(context, flatIndex);
                    localResult = applyFlat(*singleRow, flatArray, context);
                    localResult =
                            BaseVector::wrap_in_constant(rows.end(), flatIndex, localResult);
                } else {
                    localResult = applyFlat(rows, arg, context);
                }

                context.moveOrCopyResult(localResult, rows, result);
            }

        private:
            VectorPtr applyFlat(
                const SelectivityVector &rows,
                const VectorPtr &arg,
                exec::EvalCtx &context) const {
                auto array_vector = arg->as<ArrayVector>();
                POLLUX_CHECK(array_vector);
                auto elementsVector = array_vector->elements();

                auto elementsRows =
                        toElementRows(elementsVector->size(), rows, array_vector);
                exec::LocalDecodedVector elements(context, *elementsVector, elementsRows);

                vector_size_t numElements = elementsRows.end();
                vector_size_t numRows = rows.end();

                // Allocate new vectors for indices, length and offsets.
                memory::MemoryPool *pool = context.pool();
                BufferPtr newIndices = allocate_indices(numElements, pool);
                BufferPtr newSizes = allocateSizes(numRows, pool);
                BufferPtr newOffsets = allocateOffsets(numRows, pool);

                // Pointers and cursor to the raw data.
                vector_size_t indexCursor = 0;
                auto *rawIndices = newIndices->asMutable<vector_size_t>();
                auto *rawSizes = newSizes->asMutable<vector_size_t>();
                auto *rawOffsets = newOffsets->asMutable<vector_size_t>();

                // Process the rows: use a hashmap to store unique values and
                // whether it occurred once or more than once.
                melon::F14FastMap<T, bool> uniqueMap;

                rows.applyToSelected([&](vector_size_t row) {
                    auto size = array_vector->sizeAt(row);
                    auto offset = array_vector->offsetAt(row);

                    rawOffsets[row] = indexCursor;
                    vector_size_t numNulls = 0;

                    for (vector_size_t i = offset; i < offset + size; ++i) {
                        if (elements->is_null_at(i)) {
                            numNulls++;
                            if (numNulls == 2) {
                                rawIndices[indexCursor] = i;
                                indexCursor++;
                            }
                        } else {
                            T value = elements->value_at<T>(i);
                            auto it = uniqueMap.find(value);
                            if (it == uniqueMap.end()) {
                                uniqueMap[value] = true;
                            } else if (it->second) {
                                it->second = false;
                                rawIndices[indexCursor] = i;
                                indexCursor++;
                            }
                        }
                    }

                    uniqueMap.clear();
                    rawSizes[row] = indexCursor - rawOffsets[row];

                    std::sort(
                        rawIndices + rawOffsets[row],
                        rawIndices + indexCursor,
                        lib::Index2ValueNullableLess<T>(elements));
                });

                newIndices->setSize(indexCursor * sizeof(vector_size_t));
                auto newElements =
                        BaseVector::transpose(newIndices, std::move(elementsVector));

                return std::make_shared<ArrayVector>(
                    pool,
                    array_vector->type(),
                    nullptr,
                    numRows,
                    std::move(newOffsets),
                    std::move(newSizes),
                    std::move(newElements));
            }
        };

        // Validate number of parameters and types.
        void validateType(const std::vector<exec::VectorFunctionArg> &inputArgs) {
            POLLUX_USER_CHECK_EQ(
                inputArgs.size(), 1, "array_duplicates requires exactly one parameter");

            auto arrayType = inputArgs.front().type;
            POLLUX_USER_CHECK_EQ(
                arrayType->kind(),
                TypeKind::ARRAY,
                "array_duplicates requires arguments of type ARRAY");
        }

        // Create function template based on type.
        template<TypeKind kind>
        std::shared_ptr<exec::VectorFunction> createTyped(
            const std::vector<exec::VectorFunctionArg> &inputArgs) {
            POLLUX_CHECK_EQ(inputArgs.size(), 1);

            using T = typename TypeTraits<kind>::NativeType;
            return std::make_shared<ArrayDuplicatesFunction<T> >();
        }

        // Create function.
        std::shared_ptr<exec::VectorFunction> create(
            const std::string & /* name */,
            const std::vector<exec::VectorFunctionArg> &inputArgs,
            const core::QueryConfig & /*config*/) {
            validateType(inputArgs);
            auto elementType = inputArgs.front().type->childAt(0);

            return POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                        createTyped, elementType->kind(), inputArgs);
        }

        // Define function signature.
        // array(T) -> array(T) where T must be bigint or varchar.
        std::vector<std::shared_ptr<exec::FunctionSignature> > signatures() {
            return {
                exec::FunctionSignatureBuilder()
                .returnType("array(bigint)")
                .argumentType("array(bigint)")
                .build(),
                exec::FunctionSignatureBuilder()
                .returnType("array(varchar)")
                .argumentType("array(varchar)")
                .build()
            };
        }
    } // namespace

    // Register function.
    POLLUX_DECLARE_STATEFUL_VECTOR_FUNCTION(
        udf_array_duplicates,
        signatures(),
        create);
} // namespace kumo::pollux::functions
